package com.lry.basic.algorithm.union;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class Union {

    int[] par;
    int[] rank;//树的高度

    public Union(int size){
        assert size>0;
        par = new int[size];
        rank = new int[size];
        for(int i=0;i<size;i++){
            par[i] = i;
            rank[i] = 1;
        }
    }

    public Map<Integer,List<Integer>> allUnionSet(){
        findAll();
        //key是par数组的值，value是下标集合
        Map<Integer,List<Integer>> map = new ConcurrentHashMap<>();
        for(int i=0;i<par.length;i++){
            int finalI = i;
            map.computeIfAbsent(par[i], k -> new ArrayList<>()).add(finalI);
        }
        return map;
    }

    //让所有节点直连根节点
    private void findAll(){
        for(int i=0;i<par.length;i++){
            find(i);
        }
    }


    //查询树根
    public int find(int x){
        if(x < 0 || x >= par.length)
            throw new IllegalArgumentException("p is out of bound.");
        if(par[x] == x){
            return x;
        }else{
            //路径压缩
            return par[x] = find(par[x]);
        }
    }
    //合并x和y所属集合
    public void union(int x,int y){
        x = find(x);
        y = find(y);
        if(x==y)
            return;
        //rank小的向大的连
        if(rank[x]<rank[y]){
            par[x] = y;
        }else{
            par[y] = x;
            if(rank[x]==rank[y]){
               rank[x]++;
            }
        }
    }

    public boolean same(int x,int y){
        return find(x) == find(y);
    }


    public static void main(String[] args) {
        Union union = new Union(4);
        System.out.println(union.allUnionSet());
        union.union(0,1);
        System.out.println(union.allUnionSet());
        union.union(1,2);
        System.out.println(union.allUnionSet());
    }
}
